Conformal Prediction Wrappers #184

soulios-basf · 2025-07-02T11:10:26Z

Overview:
This PR introduces several new files to support conformal prediction in molecular machine learning pipelines, as well as comprehensive unit tests for pipeline and conformal prediction functionality.

conformal.py: Implements UnifiedConformalCV(single-model) and CrossConformalCV(aggregate CP) wrappers for easy conformal prediction (classification/regression) on top of scikit-learn models.
test_conformal.py: Unit tests for the conformal wrappers, covering both regression and classification.
test_pipeline.py: Unit tests for the main pipeline, including integration with conformal prediction.
advanced_04_conformal_prediction.ipynb: Example notebook showing conformal prediction on molecular data, with benchmarking and visualization.

How to test:
Run pytest on the new test files to verify correctness.
Open and run all cells in the notebook to see conformal prediction in action.

JochenSiegWork

This is the first part of my review. I mainly looked at conformal.py and only shortly on the test.

molpipeline/experimental/uncertainty/__init__.py

molpipeline/experimental/uncertainty/conformal.py

tests/test_experimental/test_uncertainty/test_conformal.py

molpipeline/experimental/uncertainty/conformal.py

…ests, random_states, hide priv functions

…nsive tests

JochenSiegWork

Now I reviewed conformal.py and the tests. In the next iteration I'll review the notebook.

The tests are well-written :) I only got some comments and requests.

JochenSiegWork · 2025-07-23T12:00:39Z

.gitignore

If you want to add these globally to the .gitignore, you shouldn't do this in this PR but create a separate one.

JochenSiegWork · 2025-07-23T12:22:01Z

molpipeline/experimental/uncertainty/conformal.py

            Self.

        Raises
        ------
        ValueError
-            If estimator_type is not 'classifier' or 'regressor'.
+            For invalid types and uninitialized.


is a word missing in this sentence?

JochenSiegWork · 2025-07-23T12:28:54Z

molpipeline/experimental/uncertainty/conformal.py


-    def predict(self, x: NDArray[Any]) -> NDArray[Any]:
+        if self.estimator_type not in {"classifier", "regressor"}:


It shouldn't be necessary to check the content of self.estimator_type again since you validated it's content in the __init__

JochenSiegWork · 2025-07-23T12:32:30Z

molpipeline/experimental/uncertainty/conformal.py

        """
+        if not self.fitted_ or self._conformal is None:
+            raise ValueError("Estimator must be fitted before calling predict")


just to make sure: It's not necessary to calibrate before predicting?

JochenSiegWork · 2025-07-23T12:48:19Z

molpipeline/experimental/uncertainty/conformal.py

+        n_bins = self.binning
+
+        def bin_func(
+            x_test: Any,


There are more specific type hints possible here. Maybe like this?

x_test: npt.NDArray[Any], model: BaseEstimator = model, y_min: float = y_min, y_max: float = y_max, n_bins: int = n_bins,

JochenSiegWork · 2025-07-24T07:28:07Z

tests/test_experimental/test_uncertainty/test_conformal.py

+            self.assertIsInstance(pred_set, list)
+            for class_idx in pred_set:
+                self.assertIsInstance(class_idx, (int, np.integer))
+                self.assertGreaterEqual(class_idx, 0)


Maybe also check <2?

JochenSiegWork · 2025-07-24T07:30:14Z

tests/test_experimental/test_uncertainty/test_conformal.py

+                self.assertIsInstance(class_idx, (int, np.integer))
+                self.assertGreaterEqual(class_idx, 0)
+
+        self.assertTrue(np.all(p_values_custom >= 0))


These bounds should also be checked int the tests that generate p-values

JochenSiegWork · 2025-07-24T07:43:41Z

tests/test_experimental/test_uncertainty/test_conformal.py

+        # Mondrian should give different results than baseline
+        self.assertFalse(np.array_equal(intervals_mondrian, intervals_baseline))
+
+    def test_cross_conformal_mondrian_both_classes(self) -> None:


I don't really understand what this test is testing. The test is called "both_classes". What does that mean and why is regression than also checked?

JochenSiegWork · 2025-07-24T07:55:21Z

molpipeline/experimental/uncertainty/conformal.py

+)
+
+
+class ConformalPredictor(BaseEstimator):  # pylint: disable=too-many-instance-attributes


We can also discuss later with Christian:

In sklearn usually you have separate classes for regression and classification, e.g. RandomForest{Regressor,Classifier}. Since there are quite some arguments and functions in the conformal classes that work for one of those but not for the other, it would actually make sense to have something like ConformalRegressor, ConformatClassifier, CrossConformalRegressor and CrossConformalClassifier. This would reduce the if/else checking for the estimator_type. But this is a more general interface decision, we can discuss with Christian later. Since we are planing to merge the conformal code into experimental it's also fine to change the interface later.

JochenSiegWork · 2025-07-24T07:59:51Z

tests/test_experimental/test_uncertainty/test_conformal.py

The test already look quite good. I think there are two more general tests missing that are necessary to ensure interoperability with the general MolPipeline workflows:

Make a test where a Pipeline is wrapped by the conformal classes. Here is an example with the CalibratedClassifierCV

MolPipeline/tests/test_pipeline.py

Line 348 in 518c763

def test_calibrated_classifier(self) -> None:

Make a test that uses the Conformal classes as part of a pipeline, like done here https://github.com/basf/neural-fingerprint-uncertainty/blob/69c4dde201d43c7d5242805e3fcb47af54d0b101/scripts/03_ml_experiments.py#L141

Make a test to check if serialization and deserialization works, when you dump the model and read it back in again. This should be done in all 3 variations with just sklearn estimator and the Pipeline variations in 1 & 2.

You can use the recursive_to_json and recursive_from_json function to test json serialization, like here

MolPipeline/tests/test_pipeline.py

Line 153 in 518c763

json_str = recursive_to_json(m_pipeline)

Analogously, you should write a test where you dump the trained conformal classes with joblib and then read them back in again and check if they are equal.

These 3 things are important to use the two conformal classes in our downstream workflows. We can also have a call to discuss the details.

soulios-basf added 2 commits June 23, 2025 17:02

add conforaml prediction from crepes along with tests and notebook

1d5a285

ruffed it

6de0e48

soulios-basf requested a review from JochenSiegWork July 2, 2025 11:10

soulios-basf and others added 10 commits July 2, 2025 13:36

mypy docstyle etc

aedc290

Merge branch 'main' into conformal_pred

4588f28

pull first

6947efe

ruffed and ready

6a1f1f0

tests ruffed

c9cab14

tests ruffed and formatted

319d282

tests rereformatted

7394ba6

fix test

abb1067

reformatted after fix

ebc6d49

reformatted after fix

a71fa5b

JochenSiegWork requested changes Jul 8, 2025

View reviewed changes

soulios-basf added 9 commits July 11, 2025 13:21

addressed pr comments, types, docus, mondrian, fit and calib flags, t…

9629327

…ests, random_states, hide priv functions

addressed comments and made mondrian and split modular and wrote exte…

f1cf6e5

…nsive tests

linted and formatted

bc239eb

moved 2 functions to utils

832e1d6

recommit moved 2 functions to utils

d2ac8fb

linters, formatters,docsig

836ba4f

removed conformal test from pipeline

bd37fcf

added ignore for too many variables in the test_conformal

a1b336d

flaked and ruffed

e8c7c26

JochenSiegWork requested changes Jul 24, 2025

View reviewed changes


		def predict(self, x: NDArray[Any]) -> NDArray[Any]:
		if self.estimator_type not in {"classifier", "regressor"}:

		)


		class ConformalPredictor(BaseEstimator): # pylint: disable=too-many-instance-attributes

Conformal Prediction Wrappers #184

Are you sure you want to change the base?

Conformal Prediction Wrappers #184

Uh oh!

Conversation

soulios-basf commented Jul 2, 2025

Uh oh!

JochenSiegWork left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

JochenSiegWork left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!